from torch.utils.data import Dataset
import os
from vae.data.base import ImagePaths


class CustomBase(Dataset):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.data = None

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        example = self.data[i]
        return example


class CustomTrain(CustomBase):
    def __init__(self, data_root, images_list_file, **kwargs):
        super().__init__()
        
        self.data = ImagePaths(data_root, images_list_file,  random_crop=True, **kwargs) 


class CustomTest(CustomBase):
    def __init__(self, data_root, images_list_file, **kwargs):
        super().__init__()
        
        self.data = ImagePaths(data_root, images_list_file, random_crop=False, **kwargs)
